#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
小米便签网页版专用导出工具
针对 https://i.mi.com/note/h5#/ 网页结构优化的导出解决方案
"""

import os
import json
import time
import logging
import argparse
import datetime
import pandas as pd
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException

# 尝试导入webdriver_manager，如果失败则提示用户手动指定驱动路径
try:
    from webdriver_manager.chrome import ChromeDriverManager
    WEBDRIVER_MANAGER_AVAILABLE = True
except ImportError:
    WEBDRIVER_MANAGER_AVAILABLE = False

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('mi_note_web_export.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class XiaomiWebNoteExporter:
    def __init__(self, export_format='excel', output_dir='./exports', headless=False, timeout=30, driver_path=None, chrome_path=None):
        """初始化小米便签网页版导出器"""
        self.export_format = export_format
        self.output_dir = output_dir
        self.timeout = timeout
        self.driver_path = driver_path
        self.chrome_path = chrome_path
        self.driver = self._init_driver(headless)
        self.wait = WebDriverWait(self.driver, self.timeout)
        
        # 创建输出目录
        os.makedirs(self.output_dir, exist_ok=True)

    def _init_driver(self, headless):
        """初始化Chrome WebDriver"""
        logger.info('初始化Chrome浏览器...')
        chrome_options = Options()
        
        # 设置中文界面
        chrome_options.add_argument('--lang=zh-CN')
        
        # 设置Chrome浏览器路径（如果提供）
        if self.chrome_path:
            logger.info(f'使用用户指定的Chrome浏览器路径: {self.chrome_path}')
            chrome_options.binary_location = self.chrome_path
        
        # 可选：设置无头模式
        if headless:
            chrome_options.add_argument('--headless')
            chrome_options.add_argument('--disable-gpu')
            chrome_options.add_argument('--window-size=1920x1080')
        
        # 禁用自动化控制特征
        chrome_options.add_experimental_option("excludeSwitches", ["enable-automation"])
        chrome_options.add_experimental_option('useAutomationExtension', False)
        
        # 初始化WebDriver
        try:
            if self.driver_path:
                # 用户指定了驱动路径
                logger.info(f'使用用户指定的ChromeDriver路径: {self.driver_path}')
                driver = webdriver.Chrome(
                    service=Service(self.driver_path),
                    options=chrome_options
                )
            else:
                # 尝试自动下载驱动
                if WEBDRIVER_MANAGER_AVAILABLE:
                    logger.info('尝试自动下载匹配的ChromeDriver...')
                    driver = webdriver.Chrome(
                        service=Service(ChromeDriverManager().install()),
                        options=chrome_options
                    )
                else:
                    raise ImportError("未找到webdriver_manager模块，请手动指定ChromeDriver路径")
            
            # 进一步隐藏自动化特征
            driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", {
                "source": "Object.defineProperty(navigator, 'webdriver', {get: () => undefined})"
            })
            
            return driver
        except Exception as e:
            error_msg = (f'初始化浏览器失败: {e}\n\n解决方案:\n'\
                       f'1. 如果是网络问题，请手动下载ChromeDriver并使用 --driver_path 参数指定路径\n'\
                       f'2. 下载与您Chrome浏览器版本匹配的ChromeDriver: https://chromedriver.chromium.org/downloads\n'\
                       f'3. 确保ChromeDriver可执行，在Windows上可能需要将其添加到系统PATH中')
            logger.error(error_msg)
            print(error_msg)  # 直接打印错误信息给用户
            raise
    
    def login(self):
        """登录小米云服务
        注意：由于验证码问题，这里需要用户手动输入验证码
        """
        logger.info('打开小米便签网页版...')
        self.driver.get('https://i.mi.com/note/h5#/')
        
        try:
            # 等待用户手动完成登录
            logger.info('请在浏览器中完成登录（包括输入验证码）...')
            logger.info('登录成功后，程序将自动继续执行...')
            
            # 等待便签列表页面加载完成
            # 这里通过判断是否存在便签列表元素来确认登录成功
            self.wait.until(
                EC.presence_of_element_located((By.CSS_SELECTOR, '.note-list-container'))
            )
            
            logger.info('登录成功，开始导出便签...')
            return True
        except Exception as e:
            logger.error(f'登录过程出错: {e}')
            return False
    
    def get_all_notes(self):
        """获取所有便签内容 - 针对网页版结构优化的实现"""
        logger.info('开始获取便签内容...')
        notes = []
        
        try:
            # 等待左侧便签列表加载完成
            note_list_container = self.wait.until(
                EC.presence_of_element_located((By.CSS_SELECTOR, '.note-list-container'))
            )
            
            # 获取所有便签项（左侧标题列表）
            # 尝试多种可能的CSS选择器
            note_selectors = ['.note-item', '.note-title-item', '.note-list-item']
            note_items = []
            
            for selector in note_selectors:
                try:
                    note_items = note_list_container.find_elements(By.CSS_SELECTOR, selector)
                    if note_items:
                        logger.info(f'使用选择器 {selector} 找到了 {len(note_items)} 个便签')
                        break
                except:
                    continue
            
            if not note_items:
                logger.warning('未找到便签列表项，尝试使用更通用的选择器...')
                # 尝试更通用的选择器
                note_items = self.driver.find_elements(By.XPATH, "//div[contains(@class, 'note') and contains(@class, 'item')]")
                
            if not note_items:
                logger.error('无法找到便签列表，请检查网页结构是否已更改')
                return notes
            
            logger.info(f'找到 {len(note_items)} 个便签')
            
            # 遍历每个便签
            for index, note_item in enumerate(note_items, 1):
                try:
                    # 滚动到当前便签项，确保可见
                    self.driver.execute_script("arguments[0].scrollIntoView();", note_item)
                    time.sleep(0.3)
                    
                    # 点击便签项打开详情
                    note_item.click()
                    time.sleep(0.8)  # 等待右侧内容加载
                    
                    # 提取便签信息
                    note_data = self._extract_note_data(index, note_item)
                    if note_data:
                        notes.append(note_data)
                        logger.info(f'已提取第 {index} 个便签: {note_data["title"] or "无标题"}')
                    
                except Exception as e:
                    logger.error(f'处理第 {index} 个便签时出错: {e}')
                    continue
            
            return notes
        except Exception as e:
            logger.error(f'获取便签内容时出错: {e}')
            return notes
    
    def _extract_note_data(self, index, note_item):
        """从当前打开的便签页面提取便签数据 - 针对左右分栏结构优化"""
        try:
            # 提取标题 - 首先尝试从左侧列表项获取
            try:
                # 尝试从左侧选中的便签项获取标题
                title = note_item.text.strip()
                if not title:
                    # 如果左侧没有，则从右侧内容区域获取
                    title_element = self.driver.find_element(By.CSS_SELECTOR, '.note-detail-title, .note-content-title, .detail-title')
                    title = title_element.text.strip()
            except:
                title = ""
            
            # 提取内容 - 从右侧内容区域获取
            content = ""
            content_selectors = ['.note-detail-content', '.note-content', '.detail-content', '.note-editor']
            
            for selector in content_selectors:
                try:
                    content_element = self.driver.find_element(By.CSS_SELECTOR, selector)
                    content = content_element.text.strip()
                    break
                except:
                    continue
            
            # 尝试提取时间信息
            time_text = ""
            try:
                time_element = self.driver.find_element(By.CSS_SELECTOR, '.note-time, .note-meta, .detail-time')
                time_text = time_element.text.strip()
            except:
                pass
            
            # 构建便签数据字典
            note_data = {
                'id': index,
                'title': title,
                'content': content,
                'time_info': time_text,
                'export_time': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
            
            return note_data
        except Exception as e:
            logger.error(f'提取便签数据时出错: {e}')
            return None
    
    def export_notes(self, notes):
        """导出便签数据"""
        if not notes:
            logger.warning('没有可导出的便签数据')
            return False
        
        logger.info(f'准备导出 {len(notes)} 个便签...')
        
        # 生成文件名
        timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        base_filename = f'mi_notes_{timestamp}'
        
        try:
            # 根据格式导出
            if self.export_format == 'excel':
                file_path = os.path.join(self.output_dir, f'{base_filename}.xlsx')
                self._export_to_excel(notes, file_path)
            elif self.export_format == 'csv':
                file_path = os.path.join(self.output_dir, f'{base_filename}.csv')
                self._export_to_csv(notes, file_path)
            elif self.export_format == 'json':
                file_path = os.path.join(self.output_dir, f'{base_filename}.json')
                self._export_to_json(notes, file_path)
            elif self.export_format == 'txt':
                file_path = os.path.join(self.output_dir, f'{base_filename}.txt')
                self._export_to_text(notes, file_path)
            
            logger.info(f'便签已成功导出到: {file_path}')
            return True
        except Exception as e:
            logger.error(f'导出便签时出错: {e}')
            return False
    
    def _export_to_excel(self, notes, file_path):
        """导出到Excel文件"""
        df = pd.DataFrame(notes)
        df.to_excel(file_path, index=False)
    
    def _export_to_csv(self, notes, file_path):
        """导出到CSV文件"""
        df = pd.DataFrame(notes)
        df.to_csv(file_path, index=False, encoding='utf-8-sig')
    
    def _export_to_json(self, notes, file_path):
        """导出到JSON文件"""
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(notes, f, ensure_ascii=False, indent=2)
    
    def _export_to_text(self, notes, file_path):
        """导出到文本文件"""
        with open(file_path, 'w', encoding='utf-8') as f:
            for note in notes:
                f.write(f"===== {note['title'] or '无标题'} =====\n")
                f.write(f"{note['content']}\n")
                f.write(f"----- {note['time_info']} -----\n\n")
    
    def run(self):
        """运行导出流程"""
        try:
            # 登录
            if not self.login():
                logger.error('登录失败，程序退出')
                return False
            
            # 获取便签
            notes = self.get_all_notes()
            
            # 导出便签
            if notes:
                self.export_notes(notes)
            else:
                logger.warning('未找到任何便签')
            
            return True
        except Exception as e:
            logger.error(f'程序运行出错: {e}')
            return False
        finally:
            # 关闭浏览器
            if self.driver:
                self.driver.quit()
                logger.info('浏览器已关闭')

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='小米便签网页版专用导出工具')
    parser.add_argument('--format', type=str, default='excel', 
                        choices=['excel', 'csv', 'json', 'txt'],
                        help='导出格式（默认：excel）')
    parser.add_argument('--output_dir', type=str, default='./exports',
                        help='输出目录（默认：./exports）')
    parser.add_argument('--headless', action='store_true',
                        help='无头模式运行（不显示浏览器窗口）')
    parser.add_argument('--timeout', type=int, default=30,
                        help='页面加载超时时间（秒，默认：30）')
    parser.add_argument('--driver_path', type=str, default=None,
                        help='手动指定ChromeDriver的路径（当自动下载失败时使用）')
    parser.add_argument('--chrome_path', type=str, default=None,
                        help='手动指定Chrome浏览器的路径（当无法自动找到Chrome时使用）')
    return parser.parse_args()

def main():
    """主函数"""
    # 解析命令行参数
    args = parse_arguments()
    
    # 显示程序信息
    logger.info('===== 小米便签网页版专用导出工具 =====')
    logger.info(f'导出格式: {args.format}')
    logger.info(f'输出目录: {args.output_dir}')
    logger.info(f'无头模式: {args.headless}')
    logger.info(f'超时时间: {args.timeout}秒')
    if args.driver_path:
        logger.info(f'ChromeDriver路径: {args.driver_path}')
    if args.chrome_path:
        logger.info(f'Chrome浏览器路径: {args.chrome_path}')
    
    # 创建导出器并运行
    exporter = XiaomiWebNoteExporter(
        export_format=args.format,
        output_dir=args.output_dir,
        headless=args.headless,
        timeout=args.timeout,
        driver_path=args.driver_path,
        chrome_path=args.chrome_path
    )
    
    exporter.run()
    
    logger.info('===== 程序执行完毕 =====')

if __name__ == '__main__':
    main()